from tasks.deep_learning.base_model import BaseScoringModel
from tensorflow.keras import layers
from tensorflow import keras


class LSTM_Model(BaseScoringModel):
    @classmethod
    def get_default_hpyer_parameters(cls):
        return {
            'layer_lstm': {
                'units': 128,
                'return_sequences': False
            },
            'layer_dense': {
                'activation': "linear"
            }
        }


    def build_model_arc(self):
        output_dim = 1
        config = self.hyper_parameters

        inputs = layers.Input(shape=(None,4))
        layer_lstm = layers.LSTM(**config['layer_lstm'])
        layer_dense = layers.Dense(output_dim,**config['layer_dense'])


        tensor = layer_lstm(inputs)
        output_tensor = layer_dense(tensor)
        self.tf_model = keras.Model(inputs, output_tensor)




